"""
图像预处理模块
负责对输入图像进行预处理，提高识别准确率
"""

import cv2
import numpy as np
from typing import Tuple, Optional
from skimage import morphology, filters


class ImageProcessor:
    """
    图像处理器类
    
    功能：
    - 图像增强
    - 噪声去除
    - 倾斜矫正
    - 尺寸标准化
    """
    
    def __init__(self, config: Optional[dict] = None):
        """
        初始化图像处理器
        
        参数:
            config (dict, optional): 配置字典
        """
        self.config = config or {}
        self.enhance = self.config.get('enhance', True)
        self.denoise = self.config.get('denoise', True)
        self.deskew = self.config.get('deskew', True)
        
    def preprocess(self, image: np.ndarray) -> np.ndarray:
        """
        对图像进行预处理
        
        参数:
            image (np.ndarray): 输入图像
            
        返回:
            np.ndarray: 处理后的图像
        """
        # 转换为灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # 图像增强
        if self.enhance:
            gray = self.enhance_image(gray)
            
        # 去噪
        if self.denoise:
            gray = self.remove_noise(gray)
            
        # 倾斜矫正
        if self.deskew:
            gray = self.deskew_image(gray)
            
        return gray
        
    def enhance_image(self, image: np.ndarray) -> np.ndarray:
        """
        增强图像对比度和亮度
        
        参数:
            image (np.ndarray): 输入灰度图像
            
        返回:
            np.ndarray: 增强后的图像
        """
        # 直方图均衡化
        enhanced = cv2.equalizeHist(image)
        
        # CLAHE（限制对比度自适应直方图均衡）
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(enhanced)
        
        return enhanced
        
    def remove_noise(self, image: np.ndarray) -> np.ndarray:
        """
        去除图像噪声
        
        参数:
            image (np.ndarray): 输入图像
            
        返回:
            np.ndarray: 去噪后的图像
        """
        # 高斯滤波
        denoised = cv2.GaussianBlur(image, (3, 3), 0)
        
        # 中值滤波
        denoised = cv2.medianBlur(denoised, 3)
        
        # 形态学操作
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        denoised = cv2.morphologyEx(denoised, cv2.MORPH_CLOSE, kernel)
        
        return denoised
        
    def deskew_image(self, image: np.ndarray) -> np.ndarray:
        """
        矫正倾斜的图像
        
        参数:
            image (np.ndarray): 输入图像
            
        返回:
            np.ndarray: 矫正后的图像
        """
        # 计算倾斜角度
        angle = self.compute_skew_angle(image)
        
        if abs(angle) < 0.5:  # 角度太小，不需要矫正
            return image
            
        # 旋转图像
        (h, w) = image.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, angle, 1.0)
        rotated = cv2.warpAffine(image, M, (w, h), 
                                flags=cv2.INTER_CUBIC, 
                                borderMode=cv2.BORDER_REPLICATE)
        
        return rotated
        
    def compute_skew_angle(self, image: np.ndarray) -> float:
        """
        计算图像倾斜角度
        
        参数:
            image (np.ndarray): 输入图像
            
        返回:
            float: 倾斜角度（度）
        """
        # 边缘检测
        edges = cv2.Canny(image, 50, 150, apertureSize=3)
        
        # 霍夫线检测
        lines = cv2.HoughLines(edges, 1, np.pi/180, 100)
        
        if lines is None:
            return 0.0
            
        # 计算所有线的角度
        angles = []
        for rho, theta in lines[:, 0]:
            angle = np.degrees(theta) - 90
            if -45 <= angle <= 45:  # 只考虑合理范围内的角度
                angles.append(angle)
                
        if not angles:
            return 0.0
            
        # 返回中位数角度
        return np.median(angles)
        
    def extract_plate_region(self, image: np.ndarray, 
                           bbox: Tuple[int, int, int, int]) -> np.ndarray:
        """
        从图像中提取车牌区域
        
        参数:
            image (np.ndarray): 原始图像
            bbox (Tuple[int, int, int, int]): 边界框坐标 (x1, y1, x2, y2)
            
        返回:
            np.ndarray: 车牌区域图像
        """
        x1, y1, x2, y2 = bbox
        
        # 添加边距
        padding = 5
        h, w = image.shape[:2]
        x1 = max(0, x1 - padding)
        y1 = max(0, y1 - padding)
        x2 = min(w, x2 + padding)
        y2 = min(h, y2 + padding)
        
        # 提取区域
        plate_region = image[y1:y2, x1:x2]
        
        return plate_region
        
    def normalize_plate_size(self, plate_image: np.ndarray,
                           target_width: int = 440,
                           target_height: int = 140) -> np.ndarray:
        """
        标准化车牌图像尺寸
        
        参数:
            plate_image (np.ndarray): 车牌图像
            target_width (int): 目标宽度
            target_height (int): 目标高度
            
        返回:
            np.ndarray: 标准化后的图像
        """
        # 保持宽高比调整大小
        h, w = plate_image.shape[:2]
        scale = min(target_width / w, target_height / h)
        
        new_width = int(w * scale)
        new_height = int(h * scale)
        
        # 调整大小
        resized = cv2.resize(plate_image, (new_width, new_height), 
                           interpolation=cv2.INTER_CUBIC)
        
        # 创建目标大小的画布
        result = np.zeros((target_height, target_width), dtype=np.uint8)
        
        # 居中放置
        y_offset = (target_height - new_height) // 2
        x_offset = (target_width - new_width) // 2
        result[y_offset:y_offset+new_height, 
               x_offset:x_offset+new_width] = resized
        
        return result


class PlateProcessor:
    """
    车牌处理器类
    
    功能：
    - 整合车牌检测、识别和图像处理
    - 提供统一的处理接口
    - 管理处理流程
    """
    
    def __init__(self, config: dict):
        """
        初始化车牌处理器
        
        参数:
            config (dict): 配置字典
        """
        self.config = config
        self.running = False
        
        # 初始化各个组件
        self.detector = None
        self.recognizer = None
        self.image_processor = None
        
        # 加载组件
        self._load_components()
        
    def _load_components(self):
        """加载各个处理组件"""
        try:
            # 导入并初始化检测器
            from .detector import PlateDetector
            self.detector = PlateDetector(self.config.get('model', {}).get('detector', {}))
            
            # 导入并初始化识别器
            from .recognizer import PlateRecognizer
            self.recognizer = PlateRecognizer(self.config.get('model', {}).get('recognizer', {}))
            
            # 初始化图像处理器
            self.image_processor = ImageProcessor(self.config.get('image', {}).get('preprocessing', {}))
            
            print("所有组件加载成功")
            
        except Exception as e:
            print(f"组件加载失败: {e}")
            raise
            
    def process_image(self, image: np.ndarray) -> list:
        """
        处理单张图像
        
        参数:
            image (np.ndarray): 输入图像
            
        返回:
            list: 处理结果列表
        """
        results = []
        
        try:
            # 检测车牌
            detections = self.detector.detect(image)
            
            # 对每个检测结果进行识别
            for detection in detections:
                bbox = detection['bbox']
                confidence = detection['confidence']
                
                # 提取车牌区域
                plate_region = self.image_processor.extract_plate_region(image, bbox)
                
                # 预处理车牌图像
                processed_plate = self.image_processor.preprocess(plate_region)
                
                # 识别车牌文字
                recognition_result = self.recognizer.recognize(processed_plate)
                
                # 组合结果
                result = {
                    'bbox': bbox,
                    'detection_confidence': confidence,
                    'plate_text': recognition_result['text'],
                    'recognition_confidence': recognition_result['confidence'],
                    'raw_text': recognition_result['raw_text']
                }
                
                results.append(result)
                
        except Exception as e:
            print(f"图像处理失败: {e}")
            
        return results
        
    def process_video_frame(self, frame: np.ndarray) -> tuple:
        """
        处理视频帧
        
        参数:
            frame (np.ndarray): 视频帧
            
        返回:
            tuple: (处理结果, 标注后的帧)
        """
        # 处理图像
        results = self.process_image(frame)
        
        # 在帧上绘制结果
        annotated_frame = self._draw_results(frame, results)
        
        return results, annotated_frame
        
    def _draw_results(self, image: np.ndarray, results: list) -> np.ndarray:
        """
        在图像上绘制检测和识别结果
        
        参数:
            image (np.ndarray): 原始图像
            results (list): 处理结果
            
        返回:
            np.ndarray: 标注后的图像
        """
        annotated = image.copy()
        
        for result in results:
            bbox = result['bbox']
            plate_text = result['plate_text']
            detection_conf = result['detection_confidence']
            recognition_conf = result['recognition_confidence']
            
            # 绘制边界框
            cv2.rectangle(annotated, 
                         (bbox[0], bbox[1]), 
                         (bbox[2], bbox[3]), 
                         (0, 255, 0), 2)
            
            # 绘制车牌文字
            if plate_text:
                label = f"{plate_text} ({recognition_conf:.2f})"
                label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
                
                # 绘制标签背景
                cv2.rectangle(annotated,
                             (bbox[0], bbox[1] - label_size[1] - 10),
                             (bbox[0] + label_size[0], bbox[1]),
                             (0, 255, 0), -1)
                
                # 绘制标签文字
                cv2.putText(annotated, label,
                           (bbox[0], bbox[1] - 5),
                           cv2.FONT_HERSHEY_SIMPLEX,
                           0.7, (0, 0, 0), 2)
                           
        return annotated
        
    def start(self):
        """启动处理器"""
        self.running = True
        
    def stop(self):
        """停止处理器"""
        self.running = False
        
    def is_running(self) -> bool:
        """检查处理器是否在运行"""
        return self.running 